import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
import os
from tqdm import tqdm
import random

def hist2d(x, y, n=100):
    # 2d histogram used in labels.png and evolve.png
    xedges, yedges = np.linspace(x.min(), x.max(), n), np.linspace(y.min(), y.max(), n)
    hist, xedges, yedges = np.histogram2d(x, y, (xedges, yedges))
    xidx = np.clip(np.digitize(x, xedges) - 1, 0, hist.shape[0] - 1)
    yidx = np.clip(np.digitize(y, yedges) - 1, 0, hist.shape[1] - 1)
    return np.log(hist[xidx, yidx])


path = r'D:\python\upload\learn_yolov4_and_yolov5\logs\\temp_truth.cache'
if os.path.isfile(path):
    label = torch.load(path)
else:
    truth = torch.load(r'D:\python\upload\learn_yolov4_and_yolov5\logs\save.cache')
    image_path = list(truth.keys())

    label = []
    for i, value in tqdm(enumerate(image_path)):
        img = cv2.imread(value)
        h, w, _ = img.shape
        labels = truth[value]
        for i in labels:
            i[0] = (i[0]+(i[2]-i[0])/2)/w
            i[1] = (i[1]+(i[3]-i[1])/2)/h
            i[2] = i[2]/w
            i[3] = i[3]/h
            label.append(i)
    label = np.array(label)
    torch.save(label,path)
a = label[:,4]
b = label[:,:4].transpose()
print(len(a))

# 样本不均衡，其实当类别的数目相差10倍才需要考虑样本均衡，20倍的话是一定要考虑的
# 处理方法一般都是直接复制小样本的数量

def labels_to_image_weights(labels, nc=2, class_weights=np.ones(80)):
    # Produces image weights based on class mAPs
    n = len(labels)
    a = np.copy(labels[:,4]).astype(np.int64)
    class_counts = np.bincount(a, minlength=nc)
    image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1)

    class_counts = np.array([np.bincount(labels[i][:, 0].astype(np.int), minlength=nc) for i in range(n)])
    image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1)
    # index = random.choices(range(n), weights=image_weights, k=1)  # weight image sample
    return image_weights

labels_to_image_weights(label)

# nc = int(a.max() + 1)
# key = np.unique(a)
# result = {}
# for k in key:
#     mask = (a == k)
#     y_new = a[mask]
#     v = y_new.size
#     result[k] = v
#
# num = 0
# weight = []
# for q in result.values():
#     weight.append(q/sum(result.values()))
#     num+=1
# print(weight)
# random.choice(range(len(label)),weight=weight,k=nc)


# fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)
# ax = ax.ravel()
# ax[0].hist(a, bins=np.linspace(0,nc,nc+1)-0.5, rwidth=0.8)
# ax[0].set_xlabel('classes')
# ax[1].scatter(b[0], b[1], c=hist2d(b[0], b[1], 90), cmap='jet')
# ax[1].set_xlabel('x')
# ax[1].set_ylabel('y')
# ax[2].scatter(b[2], b[3], c=hist2d(b[2], b[3], 90), cmap='jet')
# ax[2].set_xlabel('width')
# ax[2].set_ylabel('height')
#
# plt.savefig(os.path.dirname(path)+'\\' + 'labels.png', dpi=200)
# plt.show()
# plt.close()

